
# 根据app02.py的设计，在设置部分添加地址、模型和apikey的输入表单，点击设置按钮即可隐藏表单。在调用大模型时，使用用户输入的地址和模型。注意，地址和模型采用下拉菜单，api-key采用密码。地址和模型之间存在对饮关系，一个地址对应多个模型。请直接给出代码。
# 更新：一个地址名称对应一个url，对应一个model_name的列表。下拉选择地址名称，自动给url复制，并刷新下拉选择模型的内容


# -*- coding: utf-8 -*-
import streamlit as st
import requests
import json

st.set_page_config(
    page_title="Qwen2.5 Chat",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
)

custom_css = """
<style>
[data-testid="stAppViewContainer"] { background-color: #f5f5f5; }
[data-testid="stHeader"] { background-color: #1a73e8; }
[data-testid="stSidebar"] { background-color: #ffffff; }
.css-1d391kg { padding-top: 3rem; }
.css-1fv8s86 p { font-size: 1.1rem; }
.user-message { background-color: #1a73e8; color: white; border-radius:15px; padding:10px 15px; margin:5px 0; }
.bot-message { background-color: #e8f0fe; color: #202124; border-radius:15px; padding:10px 15px; margin:5px 0; }
</style>
"""
st.markdown(custom_css, unsafe_allow_html=True)

# 服务端配置列表，后面url补全/chat/completions
SERVER_CONFIGS = [
    {
        "name": "Ollama",
        "url": "http://localhost:11434/v1",
        "models": ["qwen2.5:3b", "deepseek-r1:1.5b","phi", "glm4"],
    },
    {
        "name": "AiGitee",
        "url": "https://ai.gitee.com/v1",
        "models": ["Qwen2.5-7B-Instruct"]
    },
    {
        "name": "ThuDS",
        "url": "https://madmodel.cs.tsinghua.edu.cn/v1",
        "models": ["DeepSeek-R1-Distill-32B", "DeepSeek-R1-671B"],
    }
]

if "settings" not in st.session_state:
    st.session_state.settings = {
        "show_form": True,
        "selected_server": SERVER_CONFIGS[0]["name"],
        "api_url": SERVER_CONFIGS[0]["url"],
        "selected_model": SERVER_CONFIGS[0]["models"][0],
        "api_key": ""
    }

if "messages" not in st.session_state:
    st.session_state.messages = []

with st.sidebar:
    st.header("设置")
    if st.button("清除聊天记录"):
        st.session_state.messages = []
    
    if st.button("显示/隐藏配置"):
        st.session_state.settings["show_form"] = not st.session_state.settings["show_form"]
    
    if st.session_state.settings["show_form"]:
        with st.form("config_form"):
            # 服务端选择
            new_server = st.selectbox(
                "服务地址",
                options=[c["name"] for c in SERVER_CONFIGS],
                index=next(i for i, c in enumerate(SERVER_CONFIGS) if c["name"] == st.session_state.settings["selected_server"])
            )
            
            # 当服务端改变时重置模型
            if new_server != st.session_state.settings["selected_server"]:
                st.session_state.settings["selected_server"] = new_server
                server_config = next(c for c in SERVER_CONFIGS if c["name"] == new_server)
                st.session_state.settings["api_url"] = server_config["url"]
                st.session_state.settings["selected_model"] = server_config["models"][0]
            
            # 模型选择
            current_models = next(c["models"] for c in SERVER_CONFIGS if c["name"] == st.session_state.settings["selected_server"])
            st.session_state.settings["selected_model"] = st.selectbox(
                "模型选择",
                options=current_models,
                index=current_models.index(st.session_state.settings["selected_model"])
            )
            
            st.session_state.settings["api_key"] = st.text_input(
                "API密钥",
                type="password",
                value=st.session_state.settings["api_key"]
            )
            
            st.form_submit_button("保存配置")

def generate_response(prompt):
    try:
        headers = {
            "Authorization": f"Bearer {st.session_state.settings['api_key']}",
            "Content-Type": "application/json"
        } if st.session_state.settings["api_key"] else {}
        
        response = requests.post(
            f"{st.session_state.settings['api_url']}/chat/completions",
            headers=headers,
            json={
                "model": st.session_state.settings["selected_model"],
                "messages": [{'role':'user', 'content': prompt}],
                "stream": False
            },
            timeout=60
        )
        return json.loads(response.text)["choices"][0]["message"]["content"]
    except Exception as e:
        return f"请求出错: {str(e)}"

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(f'<div class="{"user-message" if message["role"] == "user" else "bot-message"}">{message["content"]}</div>', 
                    unsafe_allow_html=True)

if prompt := st.chat_input("输入你的消息..."):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(f'<div class="user-message">{prompt}</div>', unsafe_allow_html=True)

    with st.chat_message("assistant"):
        with st.spinner("正在思考..."):
            # st.error(f"prompt:{str(prompt)}")
            full_response = generate_response(prompt)
            st.markdown(f'<div class="bot-message">{full_response}</div>', unsafe_allow_html=True)
    st.session_state.messages.append({"role": "assistant", "content": full_response})